import os
import shutil
import json
import ast
import subprocess
import time
import sys
import importlib.util

BASE_DIR = "."

# Dynamically add root_agent to sys.path
root_agent_path = os.path.join(BASE_DIR, "root_agent")
if root_agent_path not in sys.path:
    sys.path.insert(0, root_agent_path)

# Dynamically add memory_systems to sys.path to import AgentMetadataDB
sys.path.insert(0, os.path.join(BASE_DIR, "memory_systems"))
import memory_systems.agent_metadata_db

class AgentGenerator:
    def __init__(self):
        self.source_agent_dir = os.path.join(BASE_DIR, "agents", "general_agent")
        self.agent_metadata_db = memory_systems.agent_metadata_db.AgentMetadataDB() # Initialize AgentMetadataDB
        self.llm_client = self._initialize_llm_client()

    def _initialize_llm_client(self):
        try:
            # Use importlib to load LLMClient directly from its absolute path
            llm_module_path = os.path.join(BASE_DIR, "root_agent", "llm.py")
            spec = importlib.util.spec_from_file_location("llm_module", llm_module_path)
            if spec is None or spec.loader is None:
                raise ImportError(f"Could not load spec for LLM module from {llm_module_path}")
            llm_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(llm_module)
            print("LLMClient imported successfully.")
            return llm_module.LLMInterface()
        except ImportError as e:
            print(f"Warning: Could not import LLMClient from root_agent.llm. LLM generation will not work. Error: {e}")
            return None
        except Exception as e:
            print(f"Error initializing LLMClient: {e}")
            return None

    def _generate_agent_name_with_llm(self, description, requirements):
        if not self.llm_client:
            print("LLM client not initialized. Skipping agent name generation.")
            return None

        prompt = f"""你是一个高级AI工程师，正在为超级AI助理（SAIA）项目生成一个新的Agent变体。
请根据以下描述和要求，为新的Agent生成一个简洁、有意义的英文名称（小写，使用下划线分隔单词，例如：'data_processing_agent'）。
名称应能反映Agent的核心功能或特点。

描述：{description}
要求：{requirements}

请只输出Agent的名称，不要包含任何解释性文字或Markdown格式。
"""
        try:
            response = self.llm_client.generate_content(prompt)
            agent_name = response.strip().split('\n')[0].strip()
            # Basic validation for agent name
            if agent_name and all(c.isalnum() or c == '_' for c in agent_name) and agent_name[0].isalpha():
                return agent_name
            else:
                print(f"Generated agent name '{agent_name}' is invalid. Please ensure it's lowercase, uses underscores, and starts with a letter.")
                return None
        except Exception as e:
            print(f"Error generating agent name with LLM: {e}")
            return None
 

    def _copy_and_rename_files(self, source_dir, target_dir, old_agent_name, new_agent_name):
        if os.path.exists(target_dir):
            shutil.rmtree(target_dir)
        shutil.copytree(source_dir, target_dir)

        for root, _, files in os.walk(target_dir):
            for file_name in files:
                if file_name.endswith(".pyc"): # Skip .pyc files
                    continue

                old_file_path = os.path.join(root, file_name)
                new_file_name = file_name.replace(old_agent_name, new_agent_name)
                new_file_path = os.path.join(root, new_file_name)

                if old_file_path != new_file_path:
                    os.rename(old_file_path, new_file_path)

                try:
                    with open(new_file_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                    content = content.replace(old_agent_name, new_agent_name)
                    with open(new_file_path, 'w', encoding='utf-8') as f:
                        f.write(content)
                except Exception as e:
                    print(f"Error processing file {new_file_path}: {e}")

    def _generate_code_with_llm(self, file_path, original_content, all_agent_files_content, current_file_relative_path, requirements):
        if not self.llm_client:
            print("LLM client not initialized. Skipping code generation.")
            return original_content

        # Prepare reference files content, ensuring current file is last
        reference_files_str = ""
        # Add all other files first
        for filename, content in all_agent_files_content.items():
            if filename != current_file_relative_path:
                reference_files_str += f"{filename}\n:<file-{filename}>\n{content}\n</file-{filename}>\n\n"
        # Add the current file last
        reference_files_str += f"\n\n** 请只重写此当前文件 **: {current_file_relative_path}\n:<file-{current_file_relative_path}>\n{original_content}\n</file-{current_file_relative_path}>\n\n"
 
        prompt = f"""你是一个高级AI工程师，正在为超级AI助理（SAIA）项目生成一个新的Agent变体。
你的任务是根据提供的原始代码、项目参考文档和具体要求，生成修改后的Python代码。
请严格遵循以下规则：
1.  **适度修改运行流程、逻辑：** 可根据新的需求，适度调整相关流程和逻辑代码、常数定义等，。
2.  **仅可修改部分逻辑和参数**：不要改动整体架构、类结构、函数名或核心功能流程。
3.  **保持文件结构和命名一致**：除了Agent 文件名称的替换，其他文件和变量命名应保持原样。可新增 变量和逻辑。
4.  **兼容中文UTF-8**：确保代码中的字符串和注释兼容中文。
5.  **输出纯Python代码**：不要包含任何解释性文字、Markdown格式（除了代码块本身）或额外信息。
6.  **确保代码可运行且符合Python语法**。

---
**参考文件：**
{reference_files_str}

**当前文件路径：** {file_path}

**新Agent的要求：**
{requirements}

请根据上述要求和参考信息，生成修改后的Python代码。
        """
        try:
            response = self.llm_client.generate_content(prompt)
            if "```python" in response and "```" in response:
                code_start = response.find("```python") + len("```python")
                code_end = response.rfind("```")
                generated_code = response[code_start:code_end].strip()
            else:
                generated_code = response.strip()
            return generated_code
        except Exception as e:
            print(f"Error generating code with LLM for {file_path}: {e}")
            return original_content

    def _syntax_check(self, file_path):
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            ast.parse(content)
            return True
        except SyntaxError as e:
            print(f"Syntax error in {file_path}: {e}")
            return False
        except Exception as e:
            print(f"Error during syntax check for {file_path}: {e}")
            return False

    def _run_simple_test(self, agent_dir, new_agent_name):
        test_script_content = f"""
import sys
import os
import time
import importlib.util

# Add the project root to sys.path for the generated script
sys.path.insert(0, r"{BASE_DIR}")

# Add the new agent's directory to sys.path
sys.path.insert(0, r"{agent_dir}")

new_agent_name_in_script = "{new_agent_name}" # Define new_agent_name here for the generated script

try:
    module_name = new_agent_name_in_script
    main_file_path = os.path.join(r"{agent_dir}", f"{{module_name}}.py")

    if not os.path.exists(main_file_path):
        print(f"Error: Main agent file not found at {{main_file_path}}")
        sys.exit(1)

    spec = importlib.util.spec_from_file_location(module_name, main_file_path)
    if spec is None:
        print(f"Error: Could not load spec for {{module_name}} from {{main_file_path}}")
        sys.exit(1)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)

    agent_class_name = module_name.replace('_', ' ').title().replace(' ', '')
    if hasattr(module, agent_class_name):
        print(f"Successfully imported and found class {{agent_class_name}} in {{module_name}}.")
        sys.exit(0)
    else:
        print(f"Error: Class {{agent_class_name}} not found in module {{module_name}}.")
        sys.exit(1)

except Exception as e:
    print(f"Test failed: {{e}}")
    sys.exit(1)
"""
        test_script_path = os.path.join(BASE_DIR, "temp", f"test_{new_agent_name}.py")
        os.makedirs(os.path.dirname(test_script_path), exist_ok=True) # Ensure temp dir exists
        with open(test_script_path, 'w', encoding='utf-8') as f:
            f.write(test_script_content)

        command = ["python", test_script_path]
        try:
            process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding='utf-8')
            start_time = time.time()
            while True:
                if process.poll() is not None:
                    break
                if time.time() - start_time > 180: # 3 minutes timeout
                    process.kill()
                    print(f"Test for {new_agent_name} timed out after 3 minutes.")
                    return False
                time.sleep(1)

            stdout, stderr = process.communicate()
            if process.returncode == 0:
                print(f"Simple test for {new_agent_name} passed.")
                return True
            else:
                print(f"Simple test for {new_agent_name} failed with exit code {process.returncode}.")
                print("STDOUT:", stdout)
                print("STDERR:", stderr)
                return False
        except Exception as e:
            print(f"Error running simple test for {new_agent_name}: {e}")
            return False
        finally:
            if os.path.exists(test_script_path):
                os.remove(test_script_path)


    def _update_agent_info(self, new_agent_name, description):
        # Use AgentMetadataDB to add or update agent info
        if not self.agent_metadata_db.get_agent_info(new_agent_name):
            self.agent_metadata_db.add_agent(new_agent_name, description)
        else:
            self.agent_metadata_db.update_agent_info(new_agent_name, description=description)

    def generate_agent(self, description: str, requirements: str):
        # LLM call to generate new_agent_name
        new_agent_name = self._generate_agent_name_with_llm(description, requirements)
        if not new_agent_name:
            print("Failed to generate agent name. Aborting.")
            return False
        old_agent_name = "general_agent"
        target_agent_dir = os.path.join(BASE_DIR, "agents", new_agent_name)

        print(f"Starting to generate new agent: {new_agent_name}")

        try:
            self._copy_and_rename_files(self.source_agent_dir, target_agent_dir, old_agent_name, new_agent_name)
            print("Files copied and renamed successfully.")
        except Exception as e:
            print(f"Error during file copy/rename: {e}")
            return False
 
        success = True
        for root, _, files in os.walk(target_agent_dir):
            for file_name in files:
                file_path = os.path.join(root, file_name)
                if not file_path.endswith(".py"):
                    continue

                print(f"Processing file: {file_path}")
                original_content = ""
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        original_content = f.read()

                    # Regenerate all_agent_files_content for each file to get the latest state
                    all_agent_files_content = {}
                    for r, _, fs in os.walk(target_agent_dir):
                        for fn in fs:
                            fp = os.path.join(r, fn)
                            if fp.endswith(".py"):
                                try:
                                    with open(fp, 'r', encoding='utf-8') as f_read:
                                        relative_path_read = os.path.relpath(fp, target_agent_dir)
                                        all_agent_files_content[relative_path_read] = f_read.read()
                                except Exception as e:
                                    print(f"Error reading file {fp} for context: {e}")
                                    # Decide how to handle this error, e.g., skip this file or return False
                                    continue # Continue to next file if reading fails

                    relative_path = os.path.relpath(file_path, target_agent_dir)
                    generated_code = self._generate_code_with_llm(
                        file_path,
                        original_content,
                        all_agent_files_content,
                        relative_path,
                        requirements
                    )

                    if generated_code:
                        with open(file_path, 'w', encoding='utf-8') as f:
                            f.write(generated_code)
                        print(f"Code generated and written for {file_name}.")
                    else:
                        print(f"No code generated for {file_name}. Keeping original content.")
                        success = False
                        continue

                    if not self._syntax_check(file_path):
                        print(f"Syntax check failed for {file_name}. Reverting to original content.")
                        with open(file_path, 'w', encoding='utf-8') as f:
                            f.write(original_content)
                        success = False
                except Exception as e:
                    print(f"Error processing {file_path}: {e}")
                    success = False

        if not success:
            print("Some files failed code generation or syntax check. Agent generation partially failed.")
            # Decide whether to clean up or leave partial results
            # For now, we'll leave it for debugging.
            # shutil.rmtree(target_agent_dir)
            return None

        print(f"Running simple test for {new_agent_name}...")
        if self._run_simple_test(target_agent_dir, new_agent_name):
            print(f"Agent {new_agent_name} successfully passed simple test.")
            self._update_agent_info(new_agent_name, description)
            print(f"Agent {new_agent_name} information updated in AgentMetadataDB.")

            # Generate tools_list.txt for the new agent
            import tools.tool_manager
            tool_manager = tools.tool_manager.ToolManager()
            all_tools_description = tool_manager.get_all_tool_descriptions()

            tool_selection_prompt = f"""你是一个高级AI工程师，正在为超级AI助理（SAIA）项目生成一个新的Agent变体。
这个Agent的名称是：{new_agent_name}
描述是：{description}
具体要求是：{requirements}

以下是所有可用的工具及其描述：
{all_tools_description}

请根据新Agent的描述和要求，从上述列表中选择最适合该Agent使用的工具。
只列出工具的名称，每个工具名称一行。不要包含任何解释性文字或Markdown格式。
"""
            try:
                if self.llm_client:
                    selected_tools_response = self.llm_client.generate_content(tool_selection_prompt)
                    selected_tools = [tool.strip() for tool in selected_tools_response.split('\n') if tool.strip()]
                    tools_list_path = os.path.join(target_agent_dir, "tools_list.txt")
                    with open(tools_list_path, 'w', encoding='utf-8') as f:
                        for tool in selected_tools:
                            f.write(f"{tool}\n")
                    print(f"Generated tools_list.txt for {new_agent_name} at {tools_list_path}")
                else:
                    print("LLM client not initialized. Skipping tools_list.txt generation.")
            except Exception as e:
                print(f"Error generating tools_list.txt for {new_agent_name}: {e}")

            return new_agent_name
        else:
            print(f"Agent {new_agent_name} failed simple test. Cleaning up generated agent directory.")
            shutil.rmtree(target_agent_dir)
            return None

# Define the tool function that will be exposed
def generate_new_agent(description: str, requirements: str) -> dict:
    """通过复制现有的“general_agent”并根据指定需求使用大型语言模型修改其代码来生成新代理。
    新代理的名称将由LLM自动生成。

    Args:
        description (str): 新代理用途和特性的简要说明。
        requirements (str): LLM应如何修改代理代码的详细要求。强调对逻辑和参数的更改，而不是体系结构。

    Returns:
        dict: 一个指示成功或失败以及消息的字典。
    """
    generator = AgentGenerator()
    result_agent_name = generator.generate_agent(description, requirements)
    if result_agent_name:
        return {"status": "success", "message": f"Agent '{result_agent_name}' generated and tested successfully."}
    else:
        return {"status": "failure", "message": f"Failed to generate agent. Check logs for details."}
